import torch
import numpy as np
from torch.utils import data


class DataloaderJP(data.Dataset):

    def __init__(self, embeddings_fname, vocab_fname, seq_len):
        self.lyrics_seq = torch.Tensor(np.load('./data/data(11149,20,20).npy', allow_pickle=True))
        # (111480*20*20)
        print('load lyrics_seq')

        self.cont_attr_seq = torch.zeros(self.lyrics_seq.size())
        print('load cont_attr_seq')
        # print(self.cont_attr_seq.shape)

        self.discrete_attr_seq = torch.tensor(np.load('./data/data(11149,20,3).npy', allow_pickle=True))
        print('load discrete_attr_seq')
        # , ,  = self.create_training_data()
        # (111480*20*3)

    def __len__(self):
        #         print(len(self.lyrics_seq))
        return len(self.lyrics_seq)

    def __getitem__(self, i):
        lyrics_seq = self.lyrics_seq[i]
        cont_val_seq = self.cont_attr_seq[i]
        discrete_val_seq = self.discrete_attr_seq[i]
        # TODO: Add noise shape as a parameter in the class
        noise_seq = torch.randn(lyrics_seq.shape)

        return lyrics_seq, cont_val_seq, discrete_val_seq, noise_seq
